use alloc::alloc::GlobalAlloc;
use alloc::alloc::Layout;

use crate::memory_manager::MemoryManager;
use crate::page::Pages;
use spin::{Mutex, MutexGuard};

pub struct Allocator(Mutex<Pages>);

impl Allocator {
    pub fn new(mut pages: Pages) -> Self {
        unsafe { MemoryManager::new(&mut pages).init(pages.size()) };
        Self(Mutex::new(pages))
    }

    pub fn acquire_memory(&self) -> MutexGuard<'_, Pages> {
        loop {
            if let Some(pages) = self.0.try_lock() {
                return pages;
            }
        }
    }
}

unsafe impl GlobalAlloc for Allocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        unsafe { MemoryManager::alloc_in_mem(&mut *self.acquire_memory(), layout) }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe { MemoryManager::dealloc_in_space(&mut *self.acquire_memory(), ptr, layout) }
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        unsafe {
            MemoryManager::realloc_in_space(&mut *self.acquire_memory(), ptr, layout, new_size)
        }
    }
}
